/*
 *  Copyright 2008-2011 NVIDIA Corporation
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

#pragma once

#include <thrust/detail/config.h>
#include <thrust/detail/tuple_meta_transform.h>
#include <thrust/detail/allocator/allocator_traits.h>
#include <thrust/iterator/zip_iterator.h>
#include <thrust/iterator/iterator_traits.h>

namespace thrust
{

namespace detail
{

namespace zip_allocator_detail
{

template<typename Function> void tuple_for_each2(thrust::null_type, thrust::null_type, Function) {}

template<typename Tuple1, typename Tuple2, typename Function>
  void tuple_for_each2(Tuple1 &t1, Tuple2 &t2, Function f)
{
  f(t1.get_head(), t2.get_head());
  tuple_for_each2(t1.get_tail(), t2.get_tail(), f);
} // end tuple_for_each2()

} // end zip_allocator_detail

template<typename AllocatorTuple>
  class zip_allocator
{
  private:
    // name the allocator tuple
    typedef AllocatorTuple allocator_tuple;

    // create tuples of these two allocator traits
    typedef typename thrust::detail::tuple_meta_transform<
      allocator_tuple,
      thrust::detail::allocator_pointer
    >::type pointer_tuple;

    typedef typename thrust::detail::tuple_meta_transform<
      allocator_tuple,
      thrust::detail::allocator_const_pointer
    >::type const_pointer_tuple;

    template<typename Size>
      struct allocate_functor
    {
      Size m_cnt;

      inline allocate_functor(Size cnt)
        : m_cnt(cnt)
      {}

      template<typename Alloc>
        typename allocator_pointer<Alloc>::type
          operator()(const Alloc &a) const
      {
        // XXX tuple_meta_transform requires the tuple to be const
        //     we should fix this, but WAR it for now
        Alloc &mutable_a = const_cast<Alloc&>(a);
        return mutable_a.allocate(m_cnt);
      } // end operator()
    };

    template<typename Size>
      struct deallocate_functor
    {
      Size m_cnt;
      
      inline deallocate_functor(Size cnt)
        : m_cnt(cnt)
      {}

      template<typename Alloc, typename Pointer>
        void operator()(const Alloc &a, Pointer p) const
      {
        // XXX tuple_meta_transform requires the tuple to be const
        //     we should fix this, but WAR it for now
        Alloc &mutable_a = const_cast<Alloc&>(a);
        mutable_a.deallocate(p, m_cnt);
      } // end operator()
    };

  public:
    // our pointer type is a zip_iterator
    typedef thrust::zip_iterator<pointer_tuple>                      pointer;

    // our const pointer type is also a zip_iterator
    typedef thrust::zip_iterator<const_pointer_tuple>                const_pointer;

    // our value type is the pointer's value_type
    typedef typename thrust::iterator_value<pointer>::type           value_type;

    // our reference is the pointer's reference
    typedef typename thrust::iterator_reference<pointer>::type       reference;

    // our const reference is the const pointer's reference
    typedef typename thrust::iterator_reference<const_pointer>::type const_reference;

    // our size_type is the first allocator's size_type
    typedef typename thrust::detail::allocator_size<
      typename thrust::tuple_element<0,allocator_tuple>::type
    >::type                                                          size_type;

    // our difference_type is the pointer's difference
    typedef typename thrust::iterator_difference<pointer>::type      difference_type;

    // XXX declare this constructor __host__ __device__ to WAR warnings about calling a __host__
    //     function from a __host__ __device__ function (tuple's constructor)
    __host__ __device__ inline zip_allocator() {}

    inline ~zip_allocator() {}

    inline zip_allocator(const zip_allocator &other)
      : m_allocator_tuple(other.m_allocator_tuple)
    {}

    inline pointer allocate(size_type cnt,
                            const_pointer = const_pointer())
    {
      if(cnt > this->max_size())
      {
        throw std::bad_alloc();
      } // end if

      return pointer(thrust::detail::tuple_host_transform<thrust::detail::allocator_pointer>(m_allocator_tuple, allocate_functor<size_type>(cnt)));
    } // end allocate()

    inline void deallocate(pointer p, size_type cnt)
    {
      // avoid passing a temporary iterator tuple to tuple_for_each2
      pointer_tuple ptr_tuple = p.get_iterator_tuple();
      thrust::detail::zip_allocator_detail::tuple_for_each2(m_allocator_tuple, ptr_tuple, deallocate_functor<size_type>(cnt));
    } // end deallocate()

    inline size_type max_size() const
    {
      return thrust::get<0>(m_allocator_tuple).max_size();
    } // end max_size()

    inline bool operator==(const zip_allocator &other)
    {
      return m_allocator_tuple == other.m_allocator_tuple;
    } // end operator==()

    inline bool operator!=(const zip_allocator &other)
    {
      return m_allocator_tuple != other.m_allocator_tuple;
    } // end operator!=()

  private:
    allocator_tuple m_allocator_tuple;
};

} // end detail

} // end thrust

